package main

import (
	"fmt"
	"sync"
)

func main() {
	pool := NewPool(4)

	for i := 0; i < 100; i++ {
		taskID := i
		pool.Submit(func() {
			fmt.Printf("Task %d is being processed\n", taskID)
		})
	}

	pool.Wait()
}

type Worker struct {
	ID     int
	TaskCh chan func()
	wg     *sync.WaitGroup
}

func NewWorker(id int, wg *sync.WaitGroup) *Worker {
	return &Worker{
		ID:     id,
		TaskCh: make(chan func()),
		wg:     wg,
	}
}

func (w *Worker) Start() {
	go func() {
		for task := range w.TaskCh { // 监听channel是否有任务
			task()
		}
	}()
}

func (w *Worker) Close() {
	close(w.TaskCh)
}

// Pool 携程池
type Pool struct {
	Workers []*Worker
	wg      sync.WaitGroup
}

func NewPool(size int) *Pool {
	pool := &Pool{
		Workers: make([]*Worker, size),
	}

	// 初始化size个携程
	for i := 0; i < size; i++ {
		worker := NewWorker(i, &pool.wg) // worker和pool共用一个waitGroup
		pool.Workers[i] = worker
		worker.Start()
	}

	return pool
}

func (p *Pool) Submit(task func()) {
	p.wg.Add(1)
	worker := p.selectWorker()
	worker.TaskCh <- func() {
		task()
		p.wg.Done()
	}
}

func (p *Pool) Wait() {
	p.wg.Wait()
}

// 选择任务提交的worker
func (p *Pool) selectWorker() *Worker {
	minTasks := len(p.Workers[0].TaskCh) // 第一个携程的任务数
	selectedWorker := p.Workers[0]

	// 遍历获取任务数最少的worker
	for _, worker := range p.Workers {
		numTasks := len(worker.TaskCh) // 当前携程的任务数量
		if numTasks < minTasks {       // 当前任务数 < 第一个携程的任务数
			minTasks = numTasks
			selectedWorker = worker
		}
	}

	return selectedWorker
}
